import torch

from __00__config import Config
from __02__bert_classifer_model import BertClassifierModel

config = Config()
bert_tokenizer = config.bert_tokenizer
bert_model = BertClassifierModel()
bert_model.load_state_dict(torch.load(config.model_save_path))
bert_model.to(config.device)
bert_model.eval()

def predict_fun(data_dict):
	# 将文本取出
	text = data_dict['text']
	# 将文本转换为id
	text_token = bert_tokenizer.batch_encode_plus(
		[text],
		padding='max_length',
		max_length=config.pad_size,
		pad_to_max_length=True,
		return_tensors='pt'
	)
	inputs_ids = text_token['input_ids']
	attention_mask = text_token['attention_mask']
	# 将数据挂在到设备上
	inputs_ids = inputs_ids.to(config.device)
	attention_mask = attention_mask.to(config.device)
	# 禁用梯度
	with torch.no_grad():
		# 获取预测分数
		logits = bert_model(inputs_ids, attention_mask)
		# 获取最大预测值
		predict = torch.argmax(logits, dim=1)
		# 获取索引下标
		predict_idx = predict.item()
		# 获取对应的索引列
		predict_class = config.id2class_dict[predict_idx]
		data_dict['predict_class'] = predict_class
	return data_dict


if __name__ == '__main__':
	data_dict = {'text': '孩子已经发烧四天了，主要是晚上发烧白天还好，从傍晚开始吧！四岁宝宝反复发烧怎么办呢？那现在一直这样子反复怎么办呀？'}
	pred = predict_fun(data_dict)
	print(pred)
